package org.zhs.tree;


import java.util.NoSuchElementException;

/**
 * @Author: zhs
 * @Date: 2024/5/10 20:24
 * @Version: 1.0
 * @Description: 平衡二叉树
 * 前提： 是一个二叉搜索树
 * 所有节点：
 * （左子树高度 - 右子树高度）的绝对值 <= 1 。 其中 左子树高度 - 右子树高度 被称为平衡因子。
 *
 *
 * LL型 (节点平衡因子 2 ; 节点左子树平衡因子 1):      右旋转
 * LR型 (节点平衡因子 2 ; 节点的左子树平衡因子 -1):   左子树左旋 , 自己右旋
 * RR型 (节点平衡因子 -2 ; 节点右子树平衡因子 -1):    左旋转
 * RL型 (节点平衡因子 -2 ; 节点右子树平衡因子 1):     右子树右旋 , 自己左旋
 */
public class MyAVLTree<E extends Comparable<? super E>> {

    /**
     * 根节点
     */
    private TreeNode<E> root;

    public MyAVLTree() {
        this.root = null;
    }

    public void makeEmpty() {
        this.root = null;
    }

    public boolean isEmpty() {
        return this.root == null;
    }

    /**
     * Return the height of node, or -1, if null;
     */
    private int height(TreeNode<E> node) {
        return node == null ? -1 : node.height;
    }

    public void insert(E data) {
        root = insert(data, root);
    }

    /**
     * 添加子节点的方法
     *
     * @param data
     * @param node
     * @return
     */
    private TreeNode<E> insert(E data, TreeNode<E> node) {
        if (node == null)
            return new TreeNode<>(data, null, null);

        int compareResult = data.compareTo(node.data);
        if (compareResult < 0) {
            node.leftNode = insert(data, node.leftNode);
        } else if (compareResult > 0) {
            node.rightNode = insert(data, node.rightNode);
        } else {
            // 重复的数据
        }

        return balance(node);
    }

    private static final int ALLOWED_IMBALANCE = 1;

    /**
     * LL型 (节点平衡因子 2 ; 节点左子树平衡因子 1):      右旋转
     * LR型 (节点平衡因子 2 ; 节点的左子树平衡因子 -1):   左子树左旋 , 自己右旋
     * RR型 (节点平衡因子 -2 ; 节点右子树平衡因子 -1):    左旋转
     * RL型 (节点平衡因子 -2 ; 节点右子树平衡因子 1):     右子树右旋 , 自己左旋
     * @param node
     * @return
     */
    public TreeNode<E> balance(TreeNode<E> node) {
        if (node == null) {
            return node;
        }

        // 计算平衡因子 node.leftNode - node.rightNode = 当前节点的平衡因子
        // 这是右旋
        if (height(node.leftNode) - height(node.rightNode) > ALLOWED_IMBALANCE) {
            // 这是右旋时,新的根节点没有右子树的情况(LL型)
            if (height(node.leftNode.leftNode) >= height(root.leftNode.rightNode)) {
                node = rotateWithLeftChild(node);
            }
            // (LR型)
            else {
                node = doubleWithLeftChild(node);
            }
        }
        // 这是左旋
        else if (height(node.rightNode) - height(node.leftNode) > ALLOWED_IMBALANCE) {
            // 这是左旋时,新的根节点没有左子树的情况(RR型)
            if (height(node.rightNode.rightNode) >= height(root.rightNode.leftNode)) {
                node = rotateWithRightChild(node);
            }
            // (RL型)
            else {
                node = doubleWithRightChild(node);
            }
        }

        // 当前节点的高度 = 子树的高度 + 1
        node.height = Math.max(height(node.leftNode), height(node.rightNode)) + 1;
        return node;
    }

    /**
     * LR型 (节点平衡因子 2 ; 节点的左子树平衡因子 -1):   左子树左旋 , 自己右旋
     */
    private TreeNode<E> doubleWithLeftChild(TreeNode<E> x){
        // 先左旋左子树
        x.leftNode = rotateWithRightChild(x.leftNode);
        // 在右旋自己
        return rotateWithLeftChild(x);
    }

    /**
     * RL型 (节点平衡因子 -2 ; 节点右子树平衡因子 1):     右子树右旋 , 自己左旋
     */
    private TreeNode<E> doubleWithRightChild(TreeNode<E> x){
        // 先右旋右子树
        x.rightNode = rotateWithLeftChild(x.rightNode);
        // 再自己左旋
        return rotateWithRightChild(x);
    }


    /**
     * 右旋
     * 1.原根节点 的 左子树 变成新的根节点
     * 2.  如果新的根节点之前存在右子树,那么把它变成原根节点的左子树
     * 3.原根节点 变成 新的根节点的 右子树
     * @param x
     * @return
     */
    private TreeNode<E> rotateWithLeftChild(TreeNode<E> x) {
        // 1.第一步，保存当前节点的左子节点
        TreeNode<E> node = x.leftNode;
        // 2.左子节点更新为 当前节点的左子节点的右节点（因为该子节点的右节点一定比自己大，且父节点也比自己大，那么将该节点的右节点变成父节点的左节点）
        x.leftNode = node.rightNode;
        // 3.该子节点的右节点更新为 之前的父节点
        node.rightNode = x;

        x.height = Math.max(height(x.leftNode), height(x.rightNode)) + 1;
        node.height = Math.max(height(node.leftNode), node.height) + 1;
        return node;
    }



    /**
     * 左旋
     * 1.原根节点 的 右子树 变成新的根节点
     * 2.  如果新的根节点之前存在左子点,那么把它变成原根节点的右子树
     * 3.原根节点 变成 新的根节点的 左子树
     * @param x
     * @return
     */
    private TreeNode<E> rotateWithRightChild(TreeNode<E> x) {
        // 1.获取原根节点的右子树,准备将其设置为新的根节点
        TreeNode<E> node = x.rightNode;
        // 2.新的根节点如果有左子树,那么需要把这个左节点变成原根节点的右子树
        x.rightNode = node.leftNode;
        // 3.原根结点 变成 新根节点的左子树
        node.leftNode = x;

        x.height = Math.max(height(x.rightNode), height(x.leftNode)) + 1;
        node.height = Math.max(height(node.rightNode), node.height) + 1;
        return node;
    }

    public boolean contains(E data) {
        return contains(data, root);
    }


    /**
     * 判断树中是否包含指定的数据
     *
     * @param data
     * @param node
     * @return
     */
    private boolean contains(E data, TreeNode<E> node) {
        // 当前节点为空时返回false
        if (node == null) {
            return false;
        }

        // 传入的值和当前节点的值进行对比，判断是否相等还是从左侧或者右侧开始进行 递归
        int compareResult = data.compareTo(node.data);
        if (compareResult < 0) {
            return contains(data, node.leftNode);
        } else if (compareResult > 0) {
            return contains(data, node.rightNode);
        } else {
            return true;
        }
    }


    public E findMin() {
        return findMin(root).data;
    }

    /**
     * 查找最小的节点
     *
     * @param node
     * @return
     */
    private TreeNode<E> findMin(TreeNode<E> node) {
        if (isEmpty()) {
            throw new NoSuchElementException();
        }
        // 递归查询左侧节点
        if (node.leftNode != null) {
            return findMin(node.leftNode);
        }
        return node;
    }

    public E findMax() {
        return findMax(root).data;
    }

    /**
     * 查找最大的节点
     *
     * @param node
     * @return
     */
    private TreeNode<E> findMax(TreeNode<E> node) {
        if (isEmpty()) {
            throw new NoSuchElementException();
        }
        // 递归查询右侧节点
        if (node.rightNode != null) {
            return findMax(node.rightNode);
        }
        return node;
    }


    public void remove(E data) {
        root = remove(data, root);
    }


    /**
     * 删除节点
     *
     * @param data
     * @param node
     * @return
     */
    private TreeNode<E> remove(E data, TreeNode<E> node) {
        // 如果节点是空直接返回
        if (node == null) {
            return null;
        }

        // 然后要删除的值与当前节点的值进行对比
        int compareResult = data.compareTo(node.data);


        if (compareResult < 0) {
            // 如果比当前节点的值小，继续从左侧查找数据，直到找到对应的数据
            node.leftNode = remove(data, node.leftNode);
        } else if (compareResult > 0) {
            // 如果比当前节点的值大，继续从右侧查找数据，直到找到对应的数据
            node.rightNode = remove(data, node.rightNode);
        } else if (node.leftNode != null && node.rightNode != null) {
            // 如果与当前节点的值相等，并且找到的节点还有2个子节点
            // 那么就将找到的节点的值设置为 当前节点的右侧最小节点的值
            // 然后依次递归处理找到节点的子节点(表现上看这里还是从当前节点开始递归，实际上因为传入的值是找到节点的值，所以是从找到节点开始处理)
            node.data = findMin(node.rightNode).data;
            node.rightNode = remove(node.data, node.rightNode);
        } else {
            //如果 传入的数据 与 当前节点的值相等，并且 只有一个叶子节点 ，那么直接将当前节点设置成为这个子节点(不需要管该子节点是否还有子节点)
            node = node.leftNode != null ? node.leftNode : node.rightNode;
        }

        return balance(node);
    }


    public void prePrint() {
        prePrint(root);
    }

    /**
     * 前序遍历
     * 根 -> 左 -> 右
     * 先输出父节点、再遍历左侧子节点和右侧子节点
     *
     * @param node
     */
    private void prePrint(TreeNode<E> node) {
        // 先输出根节点
        if (node != null) {
            System.out.println(node.data);
        }
        // 左侧子节点不为空，则继续递归
        if (node.leftNode != null) {
            prePrint(node.leftNode);
        }
        // 右侧子节点不为空，则继续递归
        if (node.rightNode != null) {
            prePrint(node.rightNode);
        }
    }


    public void midPrint() {
        midPrint(root);
    }

    /**
     * 中序遍历
     * 左 -> 根 -> 右
     * 先遍历左侧子节点、再当前节点、再遍历右侧子节点
     *
     * @param node
     */
    private void midPrint(TreeNode<E> node) {
        if (node == null) {
            return;
        }
        // 先递归左侧子节点
        if (node.leftNode != null) {
            midPrint(node.leftNode);
        }

        // 然后输出当前节点
        System.out.println(node.data);

        // 再递归右侧子节点
        if (node.rightNode != null) {
            midPrint(node.rightNode);
        }
    }


    public void postPrint() {
        postPrint(root);
    }

    /**
     * 后序遍历
     * 左 -> 右 -> 根
     * 先遍历左侧子节点、再遍历右侧子节点 、再当前节点
     *
     * @param node
     */
    private void postPrint(TreeNode<E> node) {
        if (node == null) {
            return;
        }
        // 先递归左侧子节点
        if (node.leftNode != null) {
            midPrint(node.leftNode);
        }

        // 再递归右侧子节点
        if (node.rightNode != null) {
            midPrint(node.rightNode);
        }

        // 然后输出当前节点
        System.out.println(node.data);
    }


    /**
     * 树的节点
     *
     * @param <E>
     */
    private static class TreeNode<E> {
        /**
         * 数据
         */
        E data;

        /**
         * 左边节点
         */
        TreeNode<E> leftNode;

        /**
         * 右边节点
         */
        TreeNode<E> rightNode;

        int height;

        public TreeNode(E data) {
            this(data, null, null);
        }

        public TreeNode(E data, TreeNode<E> leftNode, TreeNode<E> rightNode) {
            this.data = data;
            this.leftNode = leftNode;
            this.rightNode = rightNode;
            this.height = 0;
        }
    }

}
